import cv2 as cv
import numpy as np
import gradio as gr
from PIL import Image

# 定义一个函数，用于生成调色板
def generate_palette(image):
    # 将PIL图像转换为NumPy数组
    image_np = np.array(image)
    
    # 将图像数据转换为一维数组，每个像素点的RGB值作为一个元素
    data = image_np.reshape((-1, 3))
    data = np.float32(data)

    # 设置KMeans的停止条件
    criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 10, 1.0)

    # 设置聚类的数量，这里是5
    num_clusters = 5

    # 执行KMeans聚类
    ret, labels, centers = cv.kmeans(data, num_clusters, None, criteria, 10, cv.KMEANS_RANDOM_CENTERS)

    # 将中心点的数据类型转换为uint8
    centers = np.uint8(centers)

    # 创建一个新的图像，用于显示调色板
    palette_height = 100  # 调色板的高度
    palette_width = int(100 * num_clusters)  # 调色板的宽度
    palette_image = np.zeros((palette_height + 30, palette_width, 3), dtype=np.uint8)  # 增加30像素的高度用于显示十六进制代码

    # 将每个主颜色填充到调色板中
    for i, color in enumerate(centers):
        start_col = int(palette_width / num_clusters) * i
        end_col = start_col + int(palette_width / num_clusters)
        palette_image[:100, start_col:end_col] = color

    # 将每个颜色的十六进制代码添加到调色板下方
    font = cv.FONT_HERSHEY_SIMPLEX
    for i, color in enumerate(centers):
        start_col = int(palette_width / num_clusters) * i
        end_col = start_col + int(palette_width / num_clusters)
        hex_color = "#{:02X}{:02X}{:02X}".format(color[0], color[1], color[2])
        cv.putText(palette_image, hex_color, (start_col, palette_height + 20), font, 0.5, (255, 255, 255), 1)

    # 将NumPy数组转换回PIL图像
    palette_image_pil = Image.fromarray(palette_image)

    return palette_image_pil

# 创建Gradio界面
iface = gr.Interface(
    fn=generate_palette,
    inputs=gr.inputs.Image(type="pil"),  # 使用gr.Image组件，type参数设置为"pil"以接收图片
    outputs=gr.outputs.Image(type="pil"),  # 输出也是图片
    title="KMeans图像调色板生成器",
    description="上传一张图片，生成包含五个主颜色的调色板。"
)

# 启动界面
iface.launch()
